from collections import OrderedDict
import torch
import torch.nn as nn
from models.unet import UNet

from collections import OrderedDict
import torch
import torch.nn as nn



class SiamUnet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, init_features=2):
        super(SiamUnet, self).__init__()
        self.CNN = UNet()
        features = init_features
        self.conv = nn.Conv2d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        )
        self.conv1 = nn.Conv2d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x1,x2):
        out1 = self.CNN(x1)
        out2 = self.CNN(x2)
        out = torch.cat((out1,out2),dim=1)
        out = self.conv(out)
        out = self.conv1(out)
        return torch.sigmoid(out)




